#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
配对统计检验工具 (Paired Statistical Test Tool)
版本: 1.0.1.20250130

功能描述:
    对CSV和Excel文件中的配对数据进行统计检验，包括配对T检验和Wilcoxon符号秩检验。
    根据分组列提取两组数据，按配对列排序后进行统计分析，结果以JSON格式输出。
    同时生成小提琴图直观展示两组数据的分布差异。

主要功能:
    1. 读取CSV和Excel文件
    2. 根据分组条件提取数据
    3. 执行配对T检验 (Paired t-test)
    4. 执行Wilcoxon符号秩检验 (Wilcoxon signed-rank test)
    5. 生成小提琴图可视化数据分布
    6. 输出JSON格式的统计结果

作者: AI Assistant
创建日期: 2025-01-30
更新日期: 2025-01-30 - 修复导包位置，添加小提琴图功能
"""

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Dict, Any, Tuple, Optional

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams
import seaborn as sns
from scipy import stats


def load_data_file(file_path: str):
    """
    加载CSV或Excel文件数据
    
    参数:
        file_path (str): 文件路径
        
    返回:
        pandas.DataFrame: 加载的数据框
        
    异常:
        FileNotFoundError: 文件不存在
        ValueError: 文件格式不支持或读取失败
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")
    
    file_ext = Path(file_path).suffix.lower()
    
    try:
        if file_ext == '.csv':
            return pd.read_csv(file_path)
        elif file_ext in ['.xlsx', '.xls']:
            return pd.read_excel(file_path)
        else:
            raise ValueError(f"不支持的文件格式: {file_ext}")
    except Exception as e:
        raise ValueError(f"读取文件失败: {str(e)}")


def extract_paired_data(df, pair_column: str, group_column: str, 
                       group_value1: str, group_value2: str, 
                       target_column: str) -> Tuple[list, list]:
    """
    从数据框中提取配对数据
    
    参数:
        df (pandas.DataFrame): 输入数据框
        pair_column (str): 配对列名
        group_column (str): 分组列名
        group_value1 (str): 分组值1
        group_value2 (str): 分组值2
        target_column (str): 目标值列名
        
    返回:
        Tuple[list, list]: 两组配对数据的目标值列表
        
    异常:
        ValueError: 列名不存在或数据提取失败
    """
    # 检查列是否存在
    required_columns = [pair_column, group_column, target_column]
    missing_columns = [col for col in required_columns if col not in df.columns]
    if missing_columns:
        raise ValueError(f"缺少必需的列: {missing_columns}")
    
    # 提取两组数据
    group1_data = df[df[group_column] == group_value1]
    group2_data = df[df[group_column] == group_value2]
    
    if group1_data.empty:
        raise ValueError(f"分组值1 '{group_value1}' 在列 '{group_column}' 中未找到数据")
    if group2_data.empty:
        raise ValueError(f"分组值2 '{group_value2}' 在列 '{group_column}' 中未找到数据")
    
    # 按配对列排序
    group1_sorted = group1_data.sort_values(pair_column)
    group2_sorted = group2_data.sort_values(pair_column)
    
    # 找到共同的配对标识符
    common_pairs = set(group1_sorted[pair_column]) & set(group2_sorted[pair_column])
    
    if not common_pairs:
        raise ValueError("两组数据中没有找到共同的配对标识符")
    
    # 提取配对的目标值
    group1_values = []
    group2_values = []
    
    for pair_id in sorted(common_pairs):
        group1_value = group1_sorted[group1_sorted[pair_column] == pair_id][target_column].iloc[0]
        group2_value = group2_sorted[group2_sorted[pair_column] == pair_id][target_column].iloc[0]
        
        # 检查是否为数值类型
        try:
            group1_values.append(float(group1_value))
            group2_values.append(float(group2_value))
        except (ValueError, TypeError):
            print(f"警告: 配对 {pair_id} 的数值无法转换为浮点数，已跳过")
            continue
    
    if len(group1_values) == 0:
        raise ValueError("没有有效的数值配对数据")
    
    return group1_values, group2_values


def perform_paired_ttest(group1_values: list, group2_values: list) -> Dict[str, Any]:
    """
    执行配对T检验
    
    参数:
        group1_values (list): 第一组数据
        group2_values (list): 第二组数据
        
    返回:
        Dict[str, Any]: T检验结果字典
        
    异常:
        ValueError: 数据长度不匹配或计算失败
    """
    if len(group1_values) != len(group2_values):
        raise ValueError("两组数据长度不匹配")
    
    if len(group1_values) < 2:
        raise ValueError("数据点数量不足，至少需要2个配对数据点")
    
    try:
        # 执行配对T检验
        t_statistic, p_value = stats.ttest_rel(group1_values, group2_values)
        
        # 计算描述性统计
        group1_array = np.array(group1_values)
        group2_array = np.array(group2_values)
        differences = group1_array - group2_array
        
        result = {
            "test_type": "paired_t_test",
            "test_name": "配对T检验",
            "t_statistic": float(t_statistic),
            "p_value": float(p_value),
            "degrees_of_freedom": len(group1_values) - 1,
            "sample_size": len(group1_values),
            "group1_mean": float(np.mean(group1_array)),
            "group2_mean": float(np.mean(group2_array)),
            "mean_difference": float(np.mean(differences)),
            "std_difference": float(np.std(differences, ddof=1)),
            "significant": bool(p_value < 0.05),  # 转换为Python bool类型
            "interpretation": "显著差异" if p_value < 0.05 else "无显著差异"
        }
        
        return result
        
    except Exception as e:
        raise ValueError(f"T检验计算失败: {str(e)}")


def perform_wilcoxon_test(group1_values: list, group2_values: list) -> Dict[str, Any]:
    """
    执行Wilcoxon符号秩检验
    
    参数:
        group1_values (list): 第一组数据
        group2_values (list): 第二组数据
        
    返回:
        Dict[str, Any]: Wilcoxon检验结果字典
        
    异常:
        ValueError: 数据长度不匹配或计算失败
    """
    if len(group1_values) != len(group2_values):
        raise ValueError("两组数据长度不匹配")
    
    if len(group1_values) < 2:
        raise ValueError("数据点数量不足，至少需要2个配对数据点")
    
    try:
        # 执行Wilcoxon符号秩检验
        statistic, p_value = stats.wilcoxon(group1_values, group2_values)
        
        # 计算描述性统计
        group1_array = np.array(group1_values)
        group2_array = np.array(group2_values)
        differences = group1_array - group2_array
        
        result = {
            "test_type": "wilcoxon_signed_rank_test",
            "test_name": "Wilcoxon符号秩检验",
            "statistic": float(statistic),
            "p_value": float(p_value),
            "sample_size": len(group1_values),
            "group1_median": float(np.median(group1_array)),
            "group2_median": float(np.median(group2_array)),
            "median_difference": float(np.median(differences)),
            "significant": bool(p_value < 0.05),  # 转换为Python bool类型
            "interpretation": "显著差异" if p_value < 0.05 else "无显著差异"
        }
        
        return result
        
    except Exception as e:
        raise ValueError(f"Wilcoxon检验计算失败: {str(e)}")


def perform_spearman_correlation(group1_values: list, group2_values: list) -> Dict[str, Any]:
    """
    计算两组配对数据的Spearman秩相关性系数
    
    参数:
        group1_values (list): 第一组数据
        group2_values (list): 第二组数据
        
    返回:
        Dict[str, Any]: Spearman相关性分析结果字典
        
    异常:
        ValueError: 数据长度不匹配或计算失败
    """
    if len(group1_values) != len(group2_values):
        raise ValueError("两组数据长度不匹配")
    
    if len(group1_values) < 2:
        raise ValueError("数据点数量不足，至少需要2个配对数据点")
    
    try:
        # 计算Spearman秩相关系数
        correlation_coefficient, p_value = stats.spearmanr(group1_values, group2_values)
        
        result = {
            "test_type": "spearman_rank_correlation",
            "test_name": "Spearman秩相关性分析",
            "correlation_coefficient": float(correlation_coefficient),
            "p_value": float(p_value),
            "sample_size": len(group1_values),
            "interpretation": "显著相关" if p_value < 0.05 else "无显著相关"
        }
        
        # 添加相关性强度解释
        abs_corr = abs(correlation_coefficient)
        if abs_corr >= 0.8:
            result["strength"] = "强相关"
        elif abs_corr >= 0.5:
            result["strength"] = "中等相关"
        elif abs_corr >= 0.3:
            result["strength"] = "弱相关"
        else:
            result["strength"] = "极弱相关或无相关"
            
        # 添加方向性解释
        if correlation_coefficient > 0:
            result["direction"] = "正相关"
        elif correlation_coefficient < 0:
            result["direction"] = "负相关"
        else:
            result["direction"] = "无线性关系"
        
        return result
        
    except Exception as e:
        raise ValueError(f"Spearman相关性分析计算失败: {str(e)}")


def create_violin_plot(group1_values: list, group2_values: list, 
                      group_value1: str, group_value2: str, 
                      target_column: str, output_path: Optional[str] = None) -> str:
    """
    创建小提琴图展示两组配对数据的分布
    
    参数:
        group1_values (list): 第一组数据
        group2_values (list): 第二组数据
        group_value1 (str): 第一组标签
        group_value2 (str): 第二组标签
        target_column (str): 目标列名
        output_path (Optional[str]): 输出文件路径，如果为None则自动生成
        
    返回:
        str: 生成的图片文件路径
        
    异常:
        ValueError: 数据为空或绘图失败
    """
    if not group1_values or not group2_values:
        raise ValueError("数据为空，无法生成小提琴图")
    
    try:
        rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS']
        
        # 准备数据
        data_for_plot = []
        labels = []
        
        # 添加第一组数据
        for value in group1_values:
            data_for_plot.append(value)
            labels.append(group_value1)
        
        # 添加第二组数据
        for value in group2_values:
            data_for_plot.append(value)
            labels.append(group_value2)
        
        # 创建DataFrame用于绘图
        plot_df = pd.DataFrame({
            'Group': labels,
            target_column: data_for_plot
        })
        
        # 创建图形
        plt.figure(figsize=(10, 6))
        
        # 绘制小提琴图
        sns.violinplot(data=plot_df, x='Group', y=target_column, palette='Set2')
        
        # 添加箱线图叠加显示
        sns.boxplot(data=plot_df, x='Group', y=target_column, 
                   width=0.3, boxprops=dict(alpha=0.7))
        
        # 设置标题和标签
        plt.title(f'配对数据分布对比 - {target_column}', fontsize=14, fontweight='bold')
        plt.xlabel('分组', fontsize=12)
        plt.ylabel(target_column, fontsize=12)
        
        # 添加统计信息
        group1_mean = np.mean(group1_values)
        group2_mean = np.mean(group2_values)
        group1_median = np.median(group1_values)
        group2_median = np.median(group2_values)
        
        # 在图上添加统计信息文本
        stats_text = f'{group_value1}: 均值={group1_mean:.3f}, 中位数={group1_median:.3f}\n'
        stats_text += f'{group_value2}: 均值={group2_mean:.3f}, 中位数={group2_median:.3f}'
        
        plt.text(0.02, 0.98, stats_text, transform=plt.gca().transAxes, 
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
        
        # 调整布局
        plt.tight_layout()
        
        # 确定输出路径
        if output_path is None:
            output_path = f"paired_data_violin_plot_{group_value1}_vs_{group_value2}.png"
        
        # 保存图片
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
        plt.close()  # 关闭图形以释放内存
        
        return output_path
        
    except Exception as e:
        raise ValueError(f"生成小提琴图失败: {str(e)}")


def save_results(results: Dict[str, Any], output_dir: str, file_prefix: str) -> str:
    """
    保存结果到JSON文件
    
    参数:
        results (Dict[str, Any]): 统计检验结果
        output_dir (str): 输出目录
        file_prefix (str): 文件前缀
        
    返回:
        str: 输出文件路径
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成输出文件名
    output_file = os.path.join(output_dir, f"{file_prefix}_paired_statistical_test_results.json")
    
    # 保存结果
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    return output_file


def main():
    """
    主函数 - 命令行入口点
    """
    parser = argparse.ArgumentParser(
        description="配对统计检验工具 - 对配对数据进行T检验和Wilcoxon符号秩检验，并生成小提琴图",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  paired-statistical-test data.csv --pair-column ID --group-column Treatment \\
    --group-value1 Before --group-value2 After --target-column Score \\
    --output-dir ./results --file-prefix experiment1

  paired-statistical-test data.xlsx --pair-column PatientID --group-column Phase \\
    --group-value1 Pre --group-value2 Post --target-column BloodPressure \\
    --output-dir /path/to/output --file-prefix bp_study
        """
    )
    
    parser.add_argument('input_file', help='输入的CSV或Excel文件路径')
    parser.add_argument('--pair-column', required=True, help='配对列的列名')
    parser.add_argument('--group-column', required=True, help='分组列的列名')
    parser.add_argument('--group-value1', required=True, help='分组值1')
    parser.add_argument('--group-value2', required=True, help='分组值2')
    parser.add_argument('--target-column', required=True, help='目标值列的列名')
    parser.add_argument('--output-dir', required=True, help='输出目录路径')
    parser.add_argument('--file-prefix', required=True, help='输出文件的前缀')
    parser.add_argument('--no-plot', action='store_true', help='不生成小提琴图')
    
    args = parser.parse_args()
    
    try:
        print(f"正在加载数据文件: {args.input_file}")
        df = load_data_file(args.input_file)
        print(f"成功加载数据，共 {len(df)} 行，{len(df.columns)} 列")
        
        print(f"正在提取配对数据...")
        group1_values, group2_values = extract_paired_data(
            df, args.pair_column, args.group_column,
            args.group_value1, args.group_value2, args.target_column
        )
        print(f"成功提取 {len(group1_values)} 对配对数据")
        
        print("正在执行配对T检验...")
        ttest_results = perform_paired_ttest(group1_values, group2_values)
        
        print("正在执行Wilcoxon符号秩检验...")
        wilcoxon_results = perform_wilcoxon_test(group1_values, group2_values)
        
        print("正在计算Spearman秩相关性系数...")
        spearman_results = perform_spearman_correlation(group1_values, group2_values)
        
        # 生成小提琴图（除非用户指定不生成）
        violin_plot_path = None
        if not args.no_plot:
            print("正在生成小提琴图...")
            try:
                # 确保输出目录存在
                os.makedirs(args.output_dir, exist_ok=True)
                
                # 生成小提琴图文件路径
                plot_filename = f"{args.file_prefix}_violin_plot.png"
                plot_path = os.path.join(args.output_dir, plot_filename)
                
                violin_plot_path = create_violin_plot(
                    group1_values, group2_values,
                    args.group_value1, args.group_value2,
                    args.target_column, plot_path
                )
                print(f"小提琴图已保存到: {violin_plot_path}")
            except Exception as e:
                print(f"警告: 生成小提琴图失败: {str(e)}")
                violin_plot_path = None
        
        # 组合结果
        final_results = {
            "analysis_info": {
                "input_file": args.input_file,
                "pair_column": args.pair_column,
                "group_column": args.group_column,
                "group_value1": args.group_value1,
                "group_value2": args.group_value2,
                "target_column": args.target_column,
                "paired_samples": len(group1_values),
                "violin_plot_path": violin_plot_path
            },
            "paired_t_test": ttest_results,
            "wilcoxon_signed_rank_test": wilcoxon_results,
            "spearman_rank_correlation": spearman_results
        }
        
        print("正在保存结果...")
        output_file = save_results(final_results, args.output_dir, args.file_prefix)
        
        print(f"\n=== 配对统计检验完成 ===")
        print(f"配对样本数: {len(group1_values)}")
        print(f"配对T检验 p值: {ttest_results['p_value']:.6f} ({ttest_results['interpretation']})")
        print(f"Wilcoxon检验 p值: {wilcoxon_results['p_value']:.6f} ({wilcoxon_results['interpretation']})")
        print(f"Spearman相关性系数: {spearman_results['correlation_coefficient']:.6f} ({spearman_results['interpretation']})")
        print(f"结果已保存到: {output_file}")
        if violin_plot_path:
            print(f"小提琴图已保存到: {violin_plot_path}")
        
    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()